// DEFINITIONS
#define ERROR_RETURN(retval) { fprintf(stderr, "Error %d %s:line %d: \n", retval,__FILE__,__LINE__);    exit(retval); }
#define MAXLENGTH 255
#define THREADS 256

// NBODY DEFINITIONS
#define EPSSQ 0.0025f
#define TINT 0.0125f

// INCLUDES
#include <cuda.h>
#include <papi.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <time.h>

// STRUCTS
typedef struct
{
	float mass;
	float posx, posy, posz;
	float velx, vely, velz;
	float accx, accy, accz;
} Body; 

// PROTOTYPES
void nBodyOnDevice(int, int, Body*);
__global__ void forceCalcKernel(int, Body*);
__global__ void advancingKernel(int, Body*);

int loadData(char*, int, Body*);
void prettyPrintBody(Body);
void printBody(Body);
Body processLine(char*, int);

/**
 * MAIN FUNCTION.
 */
int
main(int argc, char** argv)
{
	printf("Running %s...\n", argv[0]);

	// Seed RNG.
	srand(time(NULL));

	// Process args.
	if(argc < 4)
	{
		printf("ERRR Not enough arguments.\n");
		return -1;
	}

	// Get n.
	int n = atoi(argv[1]);

	if(n < 1)
	{
		printf("ERRR Need positive size argument: %d.\n", n);
		return -1;
	}

	// Get t.
	int t = atoi(argv[2]);

	if(t < 1)
	{
		printf("ERRR Need positive timestep argument: %d.\n", t);
		return -1;
	}

	// Get initial data.
	char* filename = argv[3];

	Body* bodies = (Body*) malloc(sizeof(Body) * n);

	int loadValue = loadData(filename, n, bodies);

	if(loadValue < 0)
	{
		printf("ERRR Could not load file: %s.\n", filename);
		return -1;
	}

	if(loadValue != n)
	{
		printf("ERRR Not enough data for n-body sim., n = %d.\n.", n);
		return -1;
	}

	// Show input.
	/*
	printf("Start positions:\n");
	for(int i = 0; i < n; i++)
	{
		printBody(bodies[i]);
	}
	*/

	// Set up PAPI.
	// Make sure PAPI library version is correct.
	char errstring[PAPI_MAX_STR_LEN];
	int retval;
	
	if((retval = PAPI_library_init(PAPI_VER_CURRENT)) != PAPI_VER_CURRENT )
	{
		fprintf(stderr, "Error: %s\n", errstring);
		return -1;
	}

	float rtime, ptime, mflops;
	long long flpins;

	// Start counting!
	//if( (retval=PAPI_start(EventSet)) != PAPI_OK)
	if( (retval=PAPI_flops(&rtime, &ptime, &flpins, &mflops)) != PAPI_OK)
		ERROR_RETURN(retval);

	// Run nbody simulation.
	nBodyOnDevice(n, t, bodies);

	// Get counts.
	if( (retval=PAPI_flops(&rtime, &ptime, &flpins, &mflops)) != PAPI_OK)
		ERROR_RETURN(retval);

	// Close PAPI.
	PAPI_shutdown();

	// Print out results...
	printf("n:\t%d\nrtime:\t%f\nptime:\t%f\n", n, rtime, ptime);

	// Show output.
	/*
	printf("End positions:\n");
	for(int i = 0; i < n; i++)
	{
		printBody(bodies[i]);
	}
	*/
	
	printf("Done!\n");

	free(bodies);
	
	return 0;
}

/**
 * GPGPU CODE.
 *
 * All this stuff is related to CUDA stuff.
 */
void
nBodyOnDevice(int n, int t, Body* bodies)
{
	Body* bodiesl;

	int size = sizeof(Body) * n;

	// Copy data to device.
	cudaMalloc((void**) &bodiesl, size);
	cudaMemcpy(bodiesl, bodies, size, cudaMemcpyHostToDevice);

	// Calculate how many blocks needed.
	int blocks = (n + THREADS - 1) / THREADS;

	// CALCULATE.
	for(int i = 0; i < t; i++)
	{
		forceCalcKernel<<<blocks, THREADS>>>(n, bodiesl);
		advancingKernel<<<blocks, THREADS>>>(n, bodiesl);
	}

	// Get data from device.
	cudaMemcpy(bodies, bodiesl, size, cudaMemcpyDeviceToHost);

	cudaFree(bodiesl);
}

__global__ void
forceCalcKernel(int n, Body* bodies)
{
	int i = threadIdx.x + (blockIdx.x * blockDim.x);

	// Guard for partially used blocks.
	if(i < n)
	{
		float px = bodies[i].posx;
		float py = bodies[i].posy;
		float pz = bodies[i].posz;
		
		float dx, dy, dz;
		float dsq, dinv;
		float scale;

		// BASIC FORCE CODE:
		for(int j = 0; j < n; j++)
		{
			if(i != j)
			{
				// Calculate deltas.
				dx = bodies[j].posx - px;
				dy = bodies[j].posy - py;
				dz = bodies[j].posz - pz;

				// Calculate inverse distance.
				dsq = dx*dx + (dy*dy + (dz*dz + EPSSQ));
				//dinv = 1.0 / sqrtf(dsq);
				dinv = rsqrtf(dsq);

				// Get scaled force.
				scale = bodies[j].mass * dinv * dinv * dinv;

				// Accumulate acceleration.
				bodies[i].accx = bodies[i].accx + (dx * scale);
				bodies[i].accy = bodies[i].accy + (dy * scale);
				bodies[i].accz = bodies[i].accz + (dz * scale);
			}
		}
	}
}

__global__ void
advancingKernel(int n, Body* bodies)
{
	int i = threadIdx.x + (blockIdx.x * blockDim.x);

	bodies[i].velx += bodies[i].accx * TINT / 2.0;
	bodies[i].vely += bodies[i].accy * TINT / 2.0;
	bodies[i].velz += bodies[i].accz * TINT / 2;
	bodies[i].posx += bodies[i].velx * TINT;
	bodies[i].posy += bodies[i].vely * TINT;
	bodies[i].posz += bodies[i].velz * TINT;
	bodies[i].velx += bodies[i].accx * TINT / 2.0;
	bodies[i].vely += bodies[i].accy * TINT / 2.0;
	bodies[i].velz += bodies[i].accz * TINT / 2.0;
}

/**
 * HELPER FUNCTIONS.
 *
 * These just support the main host/CUDA functions.
 */
int
loadData(char* filename, int n, Body* bodies)
{
	// Open filename.
	FILE* inputFile = fopen(filename, "r");

	if(inputFile == NULL)
	{
		printf("ERRR Could not open file: %s.\n", filename);
		return -1;
	}

	//printf("Reading %s...\n", filename);

	Body b;
	
	int count = 0;
	char line[MAXLENGTH+1]; // Null character needed.

	// Read and parse input.
	while(fgets(line, MAXLENGTH+1, inputFile) != NULL)
	{
		if(count < n)
		{
			//printf("Read line %d: %s\n", count, line);

			b = processLine(line, n);
			//printBody(b);

			bodies[count] = b;

			count++;
		}
		else
		{
			break;
		}
	}

	// ROLL OUT.
	fclose(inputFile);
	
	return count;
}

void
prettyPrintBody(Body b)
{
	printf("mass:\t%f\n", b.mass);
	printf("pos:\t(%f, %f, %f)\n", b.posx, b.posy, b.posz);
	printf("vel:\t(%f, %f, %f)\n", b.velx, b.vely, b.velz);
	printf("acc:\t(%f, %f, %f)\n", b.accx, b.accy, b.accz);
}

void
printBody(Body b)
{
	
	printf("%1.2e %1.2e %1.2e %1.2e %1.2e %1.2e %1.2e %1.2e\n",
		   b.posx, b.posy, b.posz,
		   b.velx, b.vely, b.velz,
		   b.accx, b.accy, b.accz);
}

Body
processLine(char* line, int n)
{
	Body b;

	char params[9][MAXLENGTH];

	char delims[] = " ";
	char* pch = strtok(line, delims);

	for(int i = 0; i < 9; i++)
	{
		strcpy(params[i], pch);
		pch = strtok(NULL, delims);
	}

	b.mass = 1.0/n;
	b.posx = (float) atof(params[0]);
	b.posy = (float) atof(params[1]);
	b.posz = (float) atof(params[2]);
	b.velx = (float) atof(params[3]);
	b.vely = (float) atof(params[4]);
	b.velz = (float) atof(params[5]);
	b.accx = (float) atof(params[6]);
	b.accy = (float) atof(params[7]);
	b.accz = (float) atof(params[8]);

	return b;
}
